
import cv2
import numpy as np
from paddle.inference import Config
from paddle.inference import create_predictor
from paddle.vision.transforms import Normalize,Compose,Resize

transform = Compose(
    [Resize((224, 224)), Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], data_format='CHW')])
class predict_rubbish():
    def __init__(self,model_file,params_files):
        self.model_file = model_file
        self.params_file = params_files

    def resize_short(self,img, target_size):
        """ resize_short """

        resized_width = 224
        resized_height = 224
        resized = cv2.resize(img, (resized_width, resized_height))
        return resized


    def crop_image(self,img, target_size, center):
        """ crop_image """
        height, width = img.shape[:2]
        size = target_size
        if center == True:
            w_start = (width - size) / 2
            h_start = (height - size) / 2
        else:
            w_start = np.random.randint(0, width - size + 1)
            h_start = np.random.randint(0, height - size + 1)
        w_end = w_start + size
        h_end = h_start + size
        img = img[int(h_start):int(h_end), int(w_start):int(w_end), :]
        return img


    def preprocess(self,img):
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
        img = self.resize_short(img, 224)
        # img = self.crop_image(img, 224, True)
        # bgr-> rgb && hwc->chw
        img = img[:, :, ::-1].astype('float32').transpose((2, 0, 1)) / 255
        img_mean = np.array(mean).reshape((3, 1, 1))
        img_std = np.array(std).reshape((3, 1, 1))
        img -= img_mean
        img /= img_std
        return img[np.newaxis, :]


    def predict_config(self,model_file, params_file):
        # 根据预测部署的实际情况，设置Config
        config = Config()
        # 读取模型文件
        config.set_prog_file(model_file)
        config.set_params_file(params_file)
        # Config默认是使用CPU预测，若要使用GPU预测，需要手动开启，设置运行的GPU卡号和分配的初始显存。
        # config.enable_use_gpu(500, 0)
        # 可以设置开启IR优化、开启内存优化。
        config.switch_ir_optim()
        config.enable_memory_optim()
        self.predictor = create_predictor(config)
        return self.predictor


    def predict(self,image, predictor):
        img=self.preprocess(image)
        input_names = predictor.get_input_names()
        input_tensor = predictor.get_input_handle(input_names[0])
        input_tensor.reshape(img.shape)
        input_tensor.copy_from_cpu(img.copy())
        # 执行Predictor
        predictor.run()
        # 获取输出
        output_names = predictor.get_output_names()
        output_tensor = predictor.get_output_handle(output_names[0])
        output_data = output_tensor.copy_to_cpu()
        print("output_names", output_names)
        return output_data
    def post_res(self,label_dict, res):
        res = res.tolist()
        res=res[0]
        # print(type(res))
        # print(max(res))
        target_index = res.index(max(res))

        print("结果是:" + "   " + label_dict[target_index])
        print(res)
        print(target_index)
        return label_dict[target_index]


def predict_run(image):
    model_file = "inference_model/model.pdmodel"
    params_file = "inference_model/model.pdiparams"
    label_dict = {0: "纸板——cardboard", 1: "玻璃——glass", 2: "金属——metal", 3: "纸张——paper",4:"塑料——plastic",5:"其他垃圾——trash",6:"水——water"}
    prog=predict_rubbish(model_file,params_file)
    # image_path="C:/Users/86167/Desktop/paper513.jpg"
    # image = cv2.imread(image_path)
    predictor = prog.predict_config(model_file, params_file)
    res = prog.predict(image, predictor)
    result=prog.post_res(label_dict, res)
    return result


